from imutils.video import VideoStream  # 添加imutils中的视频处理模块
from imutils.video import FPS  # 添加imutils中的帧率计算模块
import argparse  # 添加命令行解析模块
import imutils  
import time  
import cv2  # 添加opencv模块

# 设置命令行解析参数
ap = argparse.ArgumentParser()
ap.add_argument("-v","--video",type=str,help="path to input video file")
ap.add_argument("-t","--tracker",type=str,default="kcf",help="OpenCV object tracker type：csrt、kcf、boosting、mil、tld、medianflow、mosse")
args = vars(ap.parse_args())


OPENCV_OBJECT_TRACKERS = {
    "csrt":cv2.TrackerCSRT_create,
    "kcf":cv2.TrackerKCF_create,
    "boosting":cv2.TrackerBoosting_create,
    "mil":cv2.TrackerMIL_create,
    "tld":cv2.TrackerTLD_create,
    "medianflow":cv2.TrackerMedianFlow_create,
    "mosse":cv2.TrackerMOSSE_create}

# 根据参数选择跟踪器
tracker = OPENCV_OBJECT_TRACKERS[args["tracker"]]()


initBB = None  

# 若参数中无video，从网络摄像头获取视频；否则，从指定视频路径读取视频
if not args.get("video",False):
    print("[INFO] starting video stream...")
    vs.VideoStream(src=0).strat()
    time.sleep(1.0)
else:
    vs = cv2.VideoCapture(args["video"])
    fps = None

while True:
    frame = vs.read()  # 读取一帧数据
    # 如果是从视频路径读取视频，frame = frame[1]，frame[0]表示是否读取视频成功；否则是从网络摄像头读取视频，frame即为其本身
    frame = frame[1] if args.get("video",False) else frame

    # 如果不能读取到帧，说明视频结束，退出
    if frame is None:
        break

    # 对帧进行等比例放缩
    frame = imutils.resize(frame,width=500)
    (H,W) = frame.shape[:2]

    if initBB is not None:  # 若initBB中含有初始化数据
        (success,box) = tracker.update(frame)  # 从当前帧中预测目标位置

        if success:  # 若预测成功
            (x,y,w,h) = [int(v) for v in box]
            cv2.rectangle(frame,(x,y),(x+w,y+h),(0,255,0),2)  # 将目标框起来
        else:
            print('跟踪目标丢失！')

        fps.update()  
        fps.stop()

        # 在帧上显示信息
        info = [
            ("Tracker",args["tracker"]),
            ("Success","Yes" if success else "No"),
            ("FPS","{:.2f}".format(fps.fps()))]

        for (i,(k,v)) in enumerate(info):
            text = "{}:{}".format(k,v)
            cv2.putText(frame,text,(10,H-(i*20)+20),cv2.FONT_HERSHEY_SIMPLEX,0.6,(0,0,255),2)

    # 显示帧并延迟1ms
    cv2.imshow("Frame",frame)
    key = cv2.waitKey(1) & 0xff

    # 若按下按键‘s’    
    if key == ord("s"):
        initBB = cv2.selectROI("Frame",frame,fromCenter=False,showCrosshair=True)  # 用鼠标从帧上选取一个区域，并保存在initBB
        tracker.init(frame,initBB)  # 用initBB中保存的区域初始化跟踪器
        fps = FPS().start()  # 开始计算帧率
    # 若按下按键‘q’，退出
    elif key == ord("q"):
        break


if not args.get("video",False):  # 若从网络摄像头读取视频
    vs.stop()  # 停止读取视频
else:
    vs.release()  # 否则释放视频文件指针

cv2.destroyAllWindows()  # 关闭所有窗口
